# 按功能模块化导入库
import pandas as pd
import statsmodels.api as sm
from sqlalchemy import create_engine
import matplotlib.pyplot as plt
import seaborn as sns

# 数据库配置常量
DB_CONFIG = {
    'host': '127.0.0.1',
    'user': 'root',
    'password': 'root',
    'database': 'tushare1',
    'port': 3306,
    'charset': 'utf8mb4'
}

# 数据库连接器（使用上下文管理器）
def create_db_connection():
    return create_engine(
        f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}"
        f"@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
        f"?charset={DB_CONFIG['charset']}"
    )

# 数据获取与合并
def fetch_merged_data():
    query = """
    SELECT d.*, m.buy_lg_vol, m.sell_lg_vol, m.buy_elg_vol, 
           m.sell_elg_vol, m.net_mf_vol, i.vol as i_vol, 
           i.closes as i_vloses
    FROM date_1 d 
    JOIN moneyflows m ON m.ts_code = d.ts_code AND d.trade_date = m.trade_date
    LEFT JOIN index_daily i ON i.trade_date = d.trade_date 
        AND i.ts_code = '399001.SZ'
    WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' 
        AND d.ts_code = '000001.SZ'
    """
    with create_db_connection().connect() as conn:
        return pd.read_sql(query, conn)

# 特征工程处理
def process_features(df):
    # 收益率计算向量化操作
    for col, base in [('zd_closes', 'closes'), 
                     ('zs_closes', 'i_vloses'),
                     ('zs_vlo', 'i_vol')]:
        df[col] = df[base].pct_change().round(2)
    
    # 严格按原始逻辑处理缺失值
    return df.dropna(subset=['zd_closes', 'zs_closes', 'zs_closes'], 
                    how='all').reset_index(drop=True)

# 特征筛选器
def select_features(df):
    exclude = {'zd_closes', 'id', 'ts_code', 'trade_date', 'the_date',
              'opens', 'high', 'low', 'closes', 'pre_closes', 'changes',
              'pct_chg', 'amount', 'vol', 'buy_elg_vol', 'i_vol',
              'i_vloses', 'zs_vlo'}
    return df.select_dtypes(include='number').columns.difference(exclude).tolist()

# 可视化模块
def generate_plots(df, model):
    plt.figure(figsize=(10, 6))
    plt.scatter(model.fittedvalues, df['zd_closes'], alpha=0.6)
    plt.title('实际值 vs 预测值')
    plt.xlabel('预测值')
    plt.ylabel('实际值')
    plt.show()
    
    # 相关热力图
    features = df[['buy_lg_vol', 'sell_lg_vol', 'sell_elg_vol', 
                  'net_mf_vol', 'zs_closes']]
    sns.heatmap(features.corr(), annot=True, cmap='coolwarm', 
               fmt='.2f', linewidths=0.5)
    plt.title('特征相关性矩阵')
    plt.show()

# 主流程
if __name__ == "__main__":
    # 数据获取
    data = fetch_merged_data()
    
    # 特征处理
    processed_data = process_features(data)
    print("处理后的数据样例：")
    print(processed_data.head(2))
    
    # 模型训练
    features = select_features(processed_data)
    X = processed_data[features]
    X = sm.add_constant(X)  # 添加截距项
    y = processed_data['zd_closes']
    
    model = sm.OLS(y, X).fit()
    
    # 结果展示
    print("\n回归分析摘要：")
    print(model.summary())
    
    # 生成可视化
    generate_plots(processed_data, model)